import os
import pandas as pd
from openai import OpenAI
from langchain.prompts import PromptTemplate
from langchain.schema import BaseOutputParser
from langchain.chains import LLMChain
from langchain_openai import ChatOpenAI
from typing import List, Dict, Tuple
import json
import re
import glob
import shutil
from dotenv import load_dotenv
import time

load_dotenv()

# 阿里云百炼配置
DASHSCOPE_CONFIG = {
    "api_key": os.getenv('DASHSCOPE_API_KEY'),
    "base_url": os.getenv('BASE_URL'),
    "model": "qwen-long"
}

# 验证API密钥
if not DASHSCOPE_CONFIG["api_key"]:
    raise ValueError("请设置DASHSCOPE_API_KEY环境变量")


class QAFilterOutputParser(BaseOutputParser):
    """解析LLM输出的筛选结果"""

    def parse(self, text: str) -> List[Dict[str, any]]:
        """解析筛选结果文本"""
        results = []

        # 预处理：提取JSON内容（处理代码块格式）
        json_text = self._extract_json_content(text)

        # 尝试解析JSON格式
        try:
            if json_text.strip().startswith('[') or json_text.strip().startswith('{'):
                data = json.loads(json_text)
                if isinstance(data, list):
                    return data
                elif isinstance(data, dict) and 'results' in data:
                    return data['results']
        except json.JSONDecodeError:
            pass

        # 解析文本格式的筛选结果
        lines = text.strip().split('\n')
        current_result = {}

        for line in lines:
            line = line.strip()
            if not line:
                continue

            # 匹配序号
            if re.match(r'^\d+[.、：:]', line):
                if current_result and 'index' in current_result:
                    results.append(current_result)
                current_result = {'index': int(re.match(r'^(\d+)', line).group(1))}
                continue

            # 匹配是否保留
            keep_match = re.search(r'保留[：:]?\s*(是|否|true|false)', line, re.IGNORECASE)
            if keep_match:
                keep_value = keep_match.group(1).lower()
                current_result['keep'] = keep_value in ['是', 'true']
                continue

            # 匹配原因
            reason_match = re.search(r'原因[：:]?\s*(.+)', line)
            if reason_match:
                current_result['reason'] = reason_match.group(1).strip()
                continue

        # 添加最后一个结果
        if current_result and 'index' in current_result:
            results.append(current_result)

        return results

    def _extract_json_content(self, text: str) -> str:
        """提取JSON内容，处理代码块格式"""
        text = text.strip()

        # 检查是否包含```json代码块
        json_block_pattern = r'```json\s*\n?([\s\S]*?)\n?```'
        json_match = re.search(json_block_pattern, text, re.IGNORECASE)

        if json_match:
            # 提取代码块中的JSON内容
            return json_match.group(1).strip()

        # 如果没有代码块，返回原文本
        return text


class QAFilter:
    """问答对筛选器"""

    def __init__(self):
        # 初始化阿里云百炼模型
        self.llm = ChatOpenAI(
            api_key=DASHSCOPE_CONFIG["api_key"],
            base_url=DASHSCOPE_CONFIG["base_url"],
            model=DASHSCOPE_CONFIG["model"],
            temperature=0.3
        )

        # 创建筛选提示模板
        self.filter_prompt_template = PromptTemplate(
            input_variables=["qa_pairs"],
            template="""你是一个专业的问答对质量评估专家。请评估以下问答对的质量，判断每个问答对是否应该保留。

评估标准：
1. 问题是否清晰、具体、有意义
2. 答案是否完整、准确、与问题相关
3. 问答对是否包含有价值的信息
4. 问题和答案是否都不为空或无意义的内容

需要筛选掉的问答对类型：
- 问题或答案为空、过短或无意义
- 问题不清晰或过于模糊
- 答案不完整或与问题不相关
- 包含乱码或格式错误的内容
- 重复或冗余的问答对

问答对列表：
{qa_pairs}

请以JSON格式输出评估结果：

{{
  "results": [
    {{
      "index": 1,
      "keep": true,
    }},
    {{
      "index": 2,
      "keep": false,
    }}
  ]
}}

注意：
1. index对应问答对的序号（从1开始）
2. keep为true表示保留，false表示筛选掉
3. 只输出JSON格式的结果
"""
        )

        # 创建筛选LLM链
        self.filter_chain = LLMChain(
            llm=self.llm,
            prompt=self.filter_prompt_template,
            output_parser=QAFilterOutputParser()
        )

    def filter_qa_pairs(self, qa_pairs: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], Dict[str, int]]:
        """筛选问答对"""
        if not qa_pairs:
            return [], {'total': 0, 'kept': 0, 'filtered': 0}

        # 准备输入文本
        qa_text = ""
        for i, qa in enumerate(qa_pairs, 1):
            question = qa.get('提问', qa.get('question', '')).strip()
            answer = qa.get('回答', qa.get('answer', '')).strip()
            qa_text += f"{i}. 问题：{question}\n   答案：{answer}\n\n"

        try:
            # 调用LLM进行筛选
            filter_results = self.filter_chain.run(qa_pairs=qa_text)
            
            # 根据筛选结果过滤问答对
            filtered_qa_pairs = []
            stats = {'total': len(qa_pairs), 'kept': 0, 'filtered': 0}
            
            # 创建索引映射
            result_map = {result.get('index', 0): result for result in filter_results if isinstance(result, dict)}
            
            for i, qa in enumerate(qa_pairs, 1):
                result = result_map.get(i, {'keep': True, 'reason': '默认保留'})
                
                if result.get('keep', True):
                    # 保留问答对，添加筛选信息
                    filtered_qa = qa.copy()
                    filtered_qa['筛选原因'] = result.get('reason', '通过筛选')
                    filtered_qa_pairs.append(filtered_qa)
                    stats['kept'] += 1
                else:
                    stats['filtered'] += 1
                    print(f"    筛选掉第{i}个问答对: {result.get('reason', '未通过筛选')}")
            
            return filtered_qa_pairs, stats
            
        except Exception as e:
            print(f"筛选问答对时出错: {e}")
            # 出错时返回原始数据
            return qa_pairs, {'total': len(qa_pairs), 'kept': len(qa_pairs), 'filtered': 0}

    def read_excel_files(self, file_paths: List[str]) -> List[Dict[str, str]]:
        """读取多个Excel文件并合并数据"""
        all_qa_pairs = []
        
        for file_path in file_paths:
            if not os.path.exists(file_path):
                print(f"文件不存在: {file_path}")
                continue
                
            try:
                df = pd.read_excel(file_path)
                print(f"  读取文件: {os.path.basename(file_path)} ({len(df)} 个问答对)")
                
                # 转换为字典列表
                for _, row in df.iterrows():
                    qa_pair = {}
                    for col in df.columns:
                        qa_pair[col] = str(row[col]) if pd.notna(row[col]) else ''
                    all_qa_pairs.append(qa_pair)
                    
            except Exception as e:
                print(f"读取文件 {file_path} 时出错: {e}")
                continue
        
        return all_qa_pairs

    def save_filtered_qa_pairs(self, qa_pairs: List[Dict[str, str]], output_file: str):
        """保存筛选后的问答对到Excel文件"""
        if not qa_pairs:
            print("没有问答对可保存")
            return
        
        # 定义不需要保留的列
        excluded_columns = {
            '分段ID', '一级标题', '二级标题', '分段内序号', 
            '表格标题', '原始数据', 'section_id', 'h1_title', 
            'h2_title', 'qa_index_in_section', 'table_title', 'raw_data', '筛选原因'
        }
        
        # 准备数据
        data = []
        for i, qa in enumerate(qa_pairs, 1):
            row = {
                '标号': i,
                '提问': qa.get('提问', qa.get('question', '')),
                '回答': qa.get('回答', qa.get('answer', '')),
            }
            
            # 只保留不在排除列表中的其他列
            for key, value in qa.items():
                if (key not in excluded_columns and 
                    key not in ['提问', '回答', 'question', 'answer'] and
                    key not in row):
                    row[key] = value
                    
            data.append(row)
        
        # 创建DataFrame并保存
        df = pd.DataFrame(data)
        df.to_excel(output_file, index=False)
        print(f"筛选结果已保存到: {output_file}")

    def process_file_group(self, base_name: str, jsonqa_dir: str, csvqa_dir: str, output_dir: str):
        """处理同名文件组"""
        print(f"\n处理文件组: {base_name}")
        
        # 构建文件路径
        jsonqa_file = os.path.join(jsonqa_dir, f"{base_name}.xlsx")
        csvqa_file = os.path.join(csvqa_dir, f"{base_name}.xlsx")
        output_file = os.path.join(output_dir, f"{base_name}.xlsx")
        
        # 收集存在的文件
        existing_files = []
        if os.path.exists(jsonqa_file):
            existing_files.append(jsonqa_file)
        if os.path.exists(csvqa_file):
            existing_files.append(csvqa_file)
        
        if not existing_files:
            print(f"未找到文件: {base_name}")
            return
        
        # 读取并合并数据
        all_qa_pairs = self.read_excel_files(existing_files)
        
        if not all_qa_pairs:
            print(f"文件组 {base_name} 中没有有效数据")
            return
        
        print(f"  合并后总计: {len(all_qa_pairs)} 个问答对")
        
        # 筛选问答对
        print(f"  开始筛选问答对...")
        filtered_qa_pairs, stats = self.filter_qa_pairs(all_qa_pairs)
        
        print(f"  筛选完成: 保留 {stats['kept']}/{stats['total']} 个问答对")
        
        # 保存筛选结果
        if filtered_qa_pairs:
            self.save_filtered_qa_pairs(filtered_qa_pairs, output_file)
        else:
            print(f"文件组 {base_name} 筛选后没有保留任何问答对")

    def clear_folder(self, folder_path: str):
        """清空指定文件夹中的所有文件"""
        if os.path.exists(folder_path):
            try:
                shutil.rmtree(folder_path)
                os.makedirs(folder_path, exist_ok=True)
                print(f"已清空文件夹: {folder_path}")
            except Exception as e:
                print(f"清空文件夹失败 {folder_path}: {e}")
        else:
            os.makedirs(folder_path, exist_ok=True)
            print(f"已创建文件夹: {folder_path}")

    def batch_filter(self):
        """批量筛选问答对"""
        print("=== 问答对筛选工具 ===")
        print(f"连接到阿里云百炼服务: {DASHSCOPE_CONFIG['base_url']}")
        print(f"使用模型: {DASHSCOPE_CONFIG['model']}")
        print()
        
        # 设置目录路径
        current_dir = os.path.dirname(os.path.abspath(__file__))
        jsonqa_dir = os.path.join(os.path.dirname(current_dir), "B_generator", "jsonQA")
        csvqa_dir = os.path.join(os.path.dirname(current_dir), "B_generator", "csvQA")
        output_dir = os.path.join(current_dir, "QA_final")
        
        print(f"JSON问答目录: {jsonqa_dir}")
        print(f"CSV问答目录: {csvqa_dir}")
        print(f"输出目录: {output_dir}")
        print()
        
        # 清空并创建输出目录
        self.clear_folder(output_dir)
        
        # 收集所有文件名
        all_files = set()
        
        # 从jsonQA目录收集文件名
        if os.path.exists(jsonqa_dir):
            for file in glob.glob(os.path.join(jsonqa_dir, "*.xlsx")):
                base_name = os.path.splitext(os.path.basename(file))[0]
                all_files.add(base_name)
        
        # 从csvQA目录收集文件名
        if os.path.exists(csvqa_dir):
            for file in glob.glob(os.path.join(csvqa_dir, "*.xlsx")):
                base_name = os.path.splitext(os.path.basename(file))[0]
                all_files.add(base_name)
        
        if not all_files:
            print("未找到任何Excel文件")
            return
        
        print(f"找到 {len(all_files)} 个文件组")
        
        # 统计信息
        total_stats = {'total': 0, 'kept': 0, 'filtered': 0}
        start_time = time.time()
        
        # 处理每个文件组
        for i, base_name in enumerate(sorted(all_files), 1):
            print(f"\n[{i}/{len(all_files)}] 处理文件组: {base_name}")
            
            try:
                self.process_file_group(base_name, jsonqa_dir, csvqa_dir, output_dir)
                
                # 统计处理结果
                output_file = os.path.join(output_dir, f"{base_name}.xlsx")
                if os.path.exists(output_file):
                    df = pd.read_excel(output_file)
                    file_count = len(df)
                    total_stats['kept'] += file_count
                    print(f"  文件组 {base_name} 最终保留 {file_count} 个问答对")
                
            except Exception as e:
                print(f"处理文件组 {base_name} 时出错: {e}")
                continue
        
        # 输出总结
        end_time = time.time()
        duration = end_time - start_time
        
        print(f"\n=== 筛选完成 ===")
        print(f"处理时间: {duration:.2f} 秒")
        print(f"处理文件组数: {len(all_files)}")
        print(f"最终保留问答对总数: {total_stats['kept']}")
        print(f"结果已保存到目录: {output_dir}")


def main():
    """主函数"""
    filter_tool = QAFilter()
    filter_tool.batch_filter()


if __name__ == "__main__":
    main()